{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "YMix6tOohQKq"
      },
      "source": [
        "##### Copyright 2019 MicroNet Challenge Authors.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "JYx17utJhOVL"
      },
      "outputs": [],
      "source": [
        "# Copyright 2019 MicroNet Challenge Authors. All Rights Reserved.\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License atte\n",
        "#\n",
        "#     http://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "# =============================================================================="
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "7fh_l-RCrg0_"
      },
      "source": [
        "# Counting the Parameters and Operations for EfficientNet\n",
        "Here is the plan\n",
        "1. Create an instance of tf.keras.Model implementation of [EfficientNet](https://arxiv.org/pdf/1905.11946.pdf) using `create_model()`. \n",
        "2. Extract the operations manually from the compiled model into the framework-agnostic API defined in `micronet_challenge.counting` using ``read_model()`\n",
        "3. Using the operations in the given list, we print total parameter count using `micronet_challenge.counting.MicroNetCounter()` class.\n",
        "\n",
        "Let's start with creating the model and running an input of ones through."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "P5p1fkA3rgL_"
      },
      "outputs": [],
      "source": [
        "# Download the official EfficientNet implementation and add an init file to\n",
        "# the EfficientNet module s.t. we can use the model builders for our counting.\n",
        "%%bash \n",
        "test -d tpu || git clone https://github.com/tensorflow/tpu tpu \u0026\u0026 mv tpu/models/official/efficientnet/* ./ \n",
        "test -d gresearch || git clone https://github.com/google-research/google-research gresearch \u0026\u0026 mv gresearch/micronet_challenge ./ "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "mhvTtB-xdi7m"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "tf.logging.set_verbosity(tf.logging.ERROR)\n",
        "\n",
        "from micronet_challenge import counting\n",
        "import efficientnet_builder\n",
        "import efficientnet_model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Ljt8JxepCeEg"
      },
      "source": [
        "## 1. Creating the Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "l6OEK-1mangm"
      },
      "outputs": [],
      "source": [
        "\"\"\"Creates and read the operations of EfficientNet instances.\n",
        "\"\"\"\n",
        "DEFAULT_INPUT_SIZES = {\n",
        "    # (width_coefficient, depth_coefficient, resolution)\n",
        "    'efficientnet-b0': 224,\n",
        "    'efficientnet-b1': 240,\n",
        "    'efficientnet-b2': 260,\n",
        "    'efficientnet-b3': 300,\n",
        "    'efficientnet-b4': 380,\n",
        "    'efficientnet-b5': 456,\n",
        "    'efficientnet-b6': 528,\n",
        "    'efficientnet-b7': 600}\n",
        "\n",
        "def create_model(model_name, input_shape=None):\n",
        "  \"\"\"Creates and reads operations from the given model.\n",
        "\n",
        "  Args:\n",
        "    model_name: str, one of the DEFAULT_INPUT_SIZES.keys()\n",
        "    input_shape: str or None, if None will be read from the dictionary.\n",
        "\n",
        "  Returns:\n",
        "    list, of operations.\n",
        "  \"\"\"\n",
        "  if input_shape is None:\n",
        "    input_size = DEFAULT_INPUT_SIZES[model_name]\n",
        "    input_shape = (1, input_size, input_size, 3)\n",
        "  blocks_args, global_params = efficientnet_builder.get_model_params(\n",
        "      model_name, None)\n",
        "\n",
        "  print('global_params= %s' % str(global_params))\n",
        "  print('blocks_args= %s' % str('\\n'.join(map(str, blocks_args))))\n",
        "  tf.reset_default_graph()\n",
        "  with tf.variable_scope(model_name):\n",
        "    model = efficientnet_model.Model(blocks_args, global_params)\n",
        "  # This will initialize the variables.\n",
        "  _ = model(tf.ones((input_shape)))\n",
        "  return model, input_shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "aKgGOZ0UCf4n"
      },
      "outputs": [],
      "source": [
        "model_name = 'efficientnet-b0'\n",
        "model, input_shape = create_model(model_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Eb-oTQkrxuPw"
      },
      "source": [
        "## 2. Extracting Operations\n",
        "- We assume 'same' padding with square images/conv kernels.\n",
        "- batchnorm scales are not counted since they can be merged. Bias added for each batch norm applied on a layer's output.\n",
        "- `f_activation` can be changed to one of the followin `relu` or `swish`. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "-mI0t7FK39mX"
      },
      "outputs": [],
      "source": [
        "#@title Reading Utils\n",
        "# assumes everything square\n",
        "# returns number of pixels for which a convolution is calculated\n",
        "\n",
        "\n",
        "def read_block(block, input_size, f_activation='swish'):\n",
        "  \"\"\"Reads the operations on a single EfficientNet block.\n",
        "\n",
        "  Args:\n",
        "    block: efficientnet_model.MBConvBlock,\n",
        "    input_shape: int, square image assumed.\n",
        "    f_activation: str or None, one of 'relu', 'swish', None.\n",
        "\n",
        "  Returns:\n",
        "    list, of operations.\n",
        "  \"\"\"\n",
        "  ops = []\n",
        "  # 1\n",
        "  l_name = '_expand_conv'\n",
        "  if hasattr(block, l_name):\n",
        "    layer = getattr(block, l_name)\n",
        "    layer_temp = counting.Conv2D(\n",
        "        input_size, layer.kernel.shape.as_list(), layer.strides, layer.padding,\n",
        "        True, f_activation)  # Use bias true since batch_norm\n",
        "    ops.append((l_name, layer_temp))\n",
        "  # 2\n",
        "  l_name = '_depthwise_conv'\n",
        "  layer = getattr(block, l_name)\n",
        "  layer_temp = counting.DepthWiseConv2D(\n",
        "      input_size, layer.weights[0].shape.as_list(), layer.strides,\n",
        "      layer.padding, True, f_activation)  # Use bias true since batch_norm\n",
        "  ops.append((l_name, layer_temp))\n",
        "  # Input size might have changed.\n",
        "  input_size = counting.get_conv_output_size(\n",
        "      image_size=input_size, filter_size=layer_temp.kernel_shape[0],\n",
        "      padding=layer_temp.padding, stride=layer_temp.strides[0])\n",
        "  # 3\n",
        "  if block._has_se:\n",
        "    se_reduce = getattr(block, '_se_reduce')\n",
        "    se_expand = getattr(block, '_se_expand')\n",
        "    # Kernel has the input features in its second dimension.\n",
        "    n_channels = se_reduce.kernel.shape.as_list()[2]\n",
        "    ops.append(('_se_reduce_mean', counting.GlobalAvg(input_size, n_channels)))\n",
        "    # input size is 1\n",
        "    layer_temp = counting.Conv2D(\n",
        "        1, se_reduce.kernel.shape.as_list(), se_reduce.strides,\n",
        "        se_reduce.padding, True, f_activation)\n",
        "    ops.append(('_se_reduce', layer_temp))\n",
        "    layer_temp = counting.Conv2D(\n",
        "        1, se_expand.kernel.shape.as_list(), se_expand.strides,\n",
        "        se_expand.padding, True, 'sigmoid')\n",
        "    ops.append(('_se_expand', layer_temp))\n",
        "    ops.append(('_se_scale', counting.Scale(input_size, n_channels)))\n",
        "\n",
        "  # 4\n",
        "  l_name = '_project_conv'\n",
        "  layer = getattr(block, l_name)\n",
        "  layer_temp = counting.Conv2D(\n",
        "      input_size, layer.kernel.shape.as_list(), layer.strides, layer.padding,\n",
        "      True, None)  # Use bias true since batch_norm, no activation\n",
        "  ops.append((l_name, layer_temp))\n",
        "\n",
        "  if (block._block_args.id_skip\n",
        "      and all(s == 1 for s in block._block_args.strides)\n",
        "      and block._block_args.input_filters == block._block_args.output_filters):\n",
        "    ops.append(('_skip_add', counting.Add(input_size, n_channels)))\n",
        "  return ops, input_size\n",
        "\n",
        "\n",
        "def read_model(model, input_shape, f_activation='swish'):\n",
        "  \"\"\"Reads the operations on a single EfficientNet block.\n",
        "\n",
        "  Args:\n",
        "    model: efficientnet_model.Model,\n",
        "    input_shape: int, square image assumed.\n",
        "    f_activation: str or None, one of 'relu', 'swish', None.\n",
        "\n",
        "  Returns:\n",
        "    list, of operations.\n",
        "  \"\"\"\n",
        "  # Ensure that the input run through model\n",
        "  _ = model(tf.ones(input_shape))\n",
        "  input_size = input_shape[1]  # Assuming square\n",
        "  ops = []\n",
        "  # 1\n",
        "  l_name = '_conv_stem'\n",
        "  layer = getattr(model, l_name)\n",
        "  layer_temp = counting.Conv2D(\n",
        "      input_size, layer.weights[0].shape.as_list(), layer.strides,\n",
        "      layer.padding, True, f_activation)  # Use bias true since batch_norm\n",
        "  ops.append((l_name, layer_temp))\n",
        "  # Input size might have changed.\n",
        "  input_size = counting.get_conv_output_size(\n",
        "      image_size=input_size, filter_size=layer_temp.kernel_shape[0],\n",
        "      padding=layer_temp.padding, stride=layer_temp.strides[0])\n",
        "\n",
        "  # Blocks\n",
        "  for idx, block in enumerate(model._blocks):\n",
        "    block_ops, input_size = read_block(block, input_size,\n",
        "                                       f_activation=f_activation)\n",
        "    ops.append(('block_%d' % idx, block_ops))\n",
        "\n",
        "  # Head\n",
        "  l_name = '_conv_head'\n",
        "  layer = getattr(model, l_name)\n",
        "  layer_temp = counting.Conv2D(\n",
        "      input_size, layer.weights[0].shape.as_list(), layer.strides,\n",
        "      layer.padding, True, f_activation)  # Use bias true since batch_norm\n",
        "  n_channels_out = layer.weights[0].shape.as_list()[-1]\n",
        "  ops.append((l_name, layer_temp))\n",
        "\n",
        "  ops.append(('_avg_pooling', counting.GlobalAvg(input_size, n_channels_out)))\n",
        "\n",
        "  l_name = '_fc'\n",
        "  layer = getattr(model, l_name)\n",
        "  ops.append(('_fc', counting.FullyConnected(\n",
        "      layer.kernel.shape.as_list(), True, None)))\n",
        "  return ops\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ITN6iAU7WRSy"
      },
      "outputs": [],
      "source": [
        "F_ACTIVATION = 'swish'\n",
        "\n",
        "all_ops = read_model(model, input_shape, f_activation=F_ACTIVATION)\n",
        "print('\\n'.join(map(str, all_ops)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "3vfIlD5kt-J-"
      },
      "source": [
        "## 3. Counting\n",
        "- Let's define some constants need for counting.\n",
        "  - `INPUT_BITS` used for the inputs of the multiplication.\n",
        "  - `ACCUMULATOR_BITS` used for the accumulator of the additions.\n",
        "  - `PARAMETER_BITS` used to store individual parameter: which is equal to `INPUT_BITS` for simplicity here.\n",
        "  - `IS_DEBUG`, if True, reports the individual operations of a single block in addition to aggregations.\n",
        "- Sparsity is applied on convolutional layers and fully connected layers and reduces number of multiplies and adds of a vector product.\n",
        "- Sparsity mask is defined as a binary mask and added to the total parameter count.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "q3LhDFKxyjyC"
      },
      "outputs": [],
      "source": [
        "# add_bits_base=32, since 32 bit adds count 1 add.\n",
        "# mul_bits_base=32, since multiplications with 32 bit input count 1 multiplication.\n",
        "counter = counting.MicroNetCounter(all_ops, add_bits_base=32, mul_bits_base=32)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "8GjhUKR7fNBi"
      },
      "outputs": [],
      "source": [
        "# Constants\n",
        "INPUT_BITS = 16\n",
        "ACCUMULATOR_BITS = 32\n",
        "PARAMETER_BITS = INPUT_BITS\n",
        "SUMMARIZE_BLOCKS = True\n",
        "\n",
        "counter.print_summary(0, PARAMETER_BITS, ACCUMULATOR_BITS, INPUT_BITS, summarize_blocks=SUMMARIZE_BLOCKS)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "0LKyxKi-x8fn"
      },
      "outputs": [],
      "source": [
        "counter.print_summary(0.1, PARAMETER_BITS, ACCUMULATOR_BITS, INPUT_BITS, summarize_blocks=SUMMARIZE_BLOCKS)\n",
        "counter.print_summary(0.5, PARAMETER_BITS, ACCUMULATOR_BITS, INPUT_BITS, summarize_blocks=SUMMARIZE_BLOCKS)\n",
        "counter.print_summary(0.9, PARAMETER_BITS, ACCUMULATOR_BITS, INPUT_BITS, summarize_blocks=SUMMARIZE_BLOCKS)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "RJiBak3s0T1i"
      },
      "outputs": [],
      "source": [
        ""
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "YMix6tOohQKq"
      ],
      "last_runtime": {
        "build_target": "",
        "kind": "local"
      },
      "name": "Efficient Net: Counting.ipynb",
      "provenance": [
        {
          "file_id": "1HVT8h8hUZKdUSn8YtMyerZb6M6F79yzh",
          "timestamp": 1565107579137
        }
      ],
      "version": "0.3.2"
    },
    "kernelspec": {
      "display_name": "Python 2",
      "name": "python2"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
